import os
from getWindow import *
from autoClick import *
from calculate_cl import *

from ultralytics import YOLO

import time
import cv2
import numpy as np
import mss

from matplotlib import rcParams
# 设置支持中文的字体
rcParams['font.family'] = 'sans-serif'
rcParams['font.sans-serif'] = ['SimHei']  # 或者使用其他支持中文的字体

# 设置置信度阈值
conf_threshold = 0.75

files_to_keep = 3

def capture_game_frame(win_pos):
    top = win_pos[0]
    left = win_pos[1]
    right = win_pos[3]
    bottom = win_pos[2]

    # 截取游戏画面
    with mss.mss() as sct:
        monitor = {"top": top, "left": left, "width": right - left, "height": bottom - top}
        screenshot = sct.grab(monitor)
        frame = cv2.cvtColor(np.array(screenshot), cv2.COLOR_BGRA2BGR)
        return frame

def infer(image, model):
    # 将图像从 BGR 转换为 RGB
    img_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # 执行推理
    results = model(img_rgb)
    # 初始化变量存储类别 0 和类别 1 的目标信息
    top_class0_center = None  # 类别 0 最上方目标的中心点
    top_class1_bottom_center = None  # 类别 1 最上方目标的底部中心点
    min_class0_y = float('inf')  # 类别 0 的最小 y 值
    min_class1_y = float('inf')  # 类别 1 的最小 y 值
    detected_classes = set()  # 用来存储检测到的类别
    # 遍历检测结果
    for detection in results[0].boxes.data:
        x1, y1, x2, y2, score, class_id = detection.tolist()
        # 仅处理置信度大于或等于 0.75 的检测
        if score >= conf_threshold:
            class_id = int(class_id)
            detected_classes.add(class_id)
            if class_id == 0:  # 类别 0
                if y1 < min_class0_y:
                    min_class0_y = y1
                    top_class0_center = ((x1 + x2) / 2, (y1 + y2) / 2)
            elif class_id == 1:  # 类别 1
                if y1 < min_class1_y:
                    min_class1_y = y1
                    top_class1_bottom_center = (x1+(x2-x1)*0.5, y1+(y2-y1)*0.897) # 实际高度与像素框选高度比值
    # 计算类别 0 的中心点与类别 1 的底部中心点之间的垂直距离
    if top_class0_center and top_class1_bottom_center:
        distance = abs(top_class0_center[1] - top_class1_bottom_center[1])
    else:
        distance = None  # 如果缺少目标，返回 None
    # 如果检测到的类别数为 1 且置信度大于阈值，则保存图片并结束程序
    if len(detected_classes) == 1:
        return distance, True  # 返回距离并指示保存图像
    return distance, False  # 否则返回距离并指示不保存图像

if __name__ == '__main__':
    # 检查游戏界面开启
    win_pos = capture_jump_game_window()
    if win_pos == -1:
        print(f"[E]未找到窗口，请确保游戏正在运行且窗口标题正确\n")
        quit()
    # 检查初始界面
    begin_location = is_page('window.png')
    if begin_location == -1:
        print(f"[E]请保证游戏窗口无遮挡且处于开始界面\n")
        quit()
    # 开始游戏
    click_mouse_at_position((win_pos[1]+begin_location[0], win_pos[0]+begin_location[1]))
    # 加载训练好的YOLO模型
    model = YOLO('./ultralytics-main/best.pt')
    # 检查是否存在保存目录
    save_dir = "captured_images"
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    time.sleep(5)  # 等待5秒
    # 主循环
    while True:
        # 获取游戏帧
        frame = capture_game_frame(win_pos)
        # 查找棋子
        chess_pos = get_chess('window.png', 'chess.png')
        # 没有棋子则游戏结束 退出程序
        if chess_pos == -1:
            print(f"[W]游戏退出\n")
            quit()
        # 打印棋子位置
        print(f"[I] Chess position found: {chess_pos}")
        distance, save_image = infer(frame, model)
        print(f"D:{distance}")
        # Always save the frame with a timestamp
        timestamp = time.strftime('%Y%m%d_%H%M%S')
        save_path = os.path.join(save_dir, f"frame_{timestamp}.png")

        # 在保存新的截图之前，确保文件夹内不超过3个文件
        files_in_dir = len([name for name in os.listdir(save_dir) if os.path.isfile(os.path.join(save_dir, name))])
        if files_in_dir >= files_to_keep:
            # 获取最早的文件名
            oldest_file = min(os.listdir(save_dir), key=lambda f: os.path.getctime(os.path.join(save_dir, f)))
            # 删除最早的文件
            os.remove(os.path.join(save_dir, oldest_file))
        
        cv2.imwrite(save_path, frame)
        print(f"[+] 截图保存到: {save_path}")
        
        if save_image:
            break  # 结束程序

        timeDelayMs = round(4.8915*distance) # 像素距离与按键时常的换算
        print(f"T:{timeDelayMs}ms")
        mouse_press(timeDelayMs)
        time.sleep(4)  # 等待4秒
